#include "MJ_interface.h"
#include "PVT_ctrl.h"
#include "data_logger.h"
#include "data_bus.h"
#include "pino_kin_dyn.h"
#include "useful_math.h"
#include "wbc_priority.h"
#include "mpc.h"
#include "gait_scheduler.h"
#include "foot_placement.h"
#include "joystick_interpreter.h"
#include <thread>
#include <mutex>
#include <mujoco/mujoco.h>


class WalkMpcWbcControllerImp
{
public:
    ~WalkMpcWbcControllerImp()
    {
        m_running = false;
        if (m_thread && m_thread->joinable())
            m_thread->join();
    }

    void Init(mjModel* model, mjData* data)
    {
        mj_model = model;
        mj_data = data;

        // initialize classes
        mj_interface = std::make_unique<MJ_Interface>(mj_model, mj_data); // data interface for Mujoco
        kinDynSolver = std::make_unique<Pin_KinDyn>("../models/AzureLoong.urdf"); // kinematics and dynamics solver
        RobotState = std::make_unique<DataBus>(kinDynSolver->model_nv); // data bus
        WBC_solv = std::make_unique<WBC_priority>(kinDynSolver->model_nv, 18, 22, 0.7, mj_model->opt.timestep); // WBC solver
        MPC_solv = std::make_unique<MPC>(dt_200Hz); // mpc controller
        gaitScheduler = std::make_unique<GaitScheduler>(0.25, mj_model->opt.timestep); // gait scheduler
        pvtCtr = std::make_unique<PVT_Ctr>(mj_model->opt.timestep, "../common/joint_ctrl_config.json"); // PVT joint control
        footPlacement = std::make_unique<FootPlacement>(); // foot-placement planner
        jsInterp = std::make_unique<JoyStickInterpreter>(mj_model->opt.timestep); // desired baselink velocity generator
        // DataLogger logger("../record/datalog.log"); // data logger

        model_nv = kinDynSolver->model_nv;

        RobotState->width_hips = 0.229;
        footPlacement->kp_vx = 0.03;
        footPlacement->kp_vy = 0.03;
        footPlacement->kp_wz = 0.03;
        footPlacement->stepHeight = 0.15;
        footPlacement->legLength = stand_legLength;

        mju_copy(mj_data->qpos, mj_model->key_qpos, mj_model->nq * 1); // set ini pos in Mujoco

        motors_pos_des.resize(model_nv - 6, 0);
        motors_pos_cur.resize(model_nv - 6, 0);
        motors_vel_des.resize(model_nv - 6, 0);
        motors_vel_cur.resize(model_nv - 6, 0);
        motors_tau_des.resize(model_nv - 6, 0);
        motors_tau_cur.resize(model_nv - 6, 0);

        // ini position and posture for foot-end and hand
        Eigen::Vector3d fe_l_pos_L_des = { -0.018, 0.113, -stand_legLength };
        Eigen::Vector3d fe_r_pos_L_des = { -0.018, -0.116, -stand_legLength };
        Eigen::Vector3d fe_l_eul_L_des = { -0.000, -0.008, -0.000 };
        Eigen::Vector3d fe_r_eul_L_des = { 0.000, -0.008, 0.000 };
        Eigen::Matrix3d fe_l_rot_des = eul2Rot(fe_l_eul_L_des(0), fe_l_eul_L_des(1), fe_l_eul_L_des(2));
        Eigen::Matrix3d fe_r_rot_des = eul2Rot(fe_r_eul_L_des(0), fe_r_eul_L_des(1), fe_r_eul_L_des(2));

        Eigen::Vector3d hd_l_pos_L_des = { -0.02, 0.32, -0.159 };
        Eigen::Vector3d hd_r_pos_L_des = { -0.02, -0.32, -0.159 };
        Eigen::Vector3d hd_l_eul_L_des = { -1.7581, 0.2129, 2.9581 };
        Eigen::Vector3d hd_r_eul_L_des = { 1.7581, 0.2129, -2.9581 };
        Eigen::Matrix3d hd_l_rot_des = eul2Rot(hd_l_eul_L_des(0), hd_l_eul_L_des(1), hd_l_eul_L_des(2));
        Eigen::Matrix3d hd_r_rot_des = eul2Rot(hd_r_eul_L_des(0), hd_r_eul_L_des(1), hd_r_eul_L_des(2));

        resLeg = kinDynSolver->computeInK_Leg(fe_l_rot_des, fe_l_pos_L_des, fe_r_rot_des, fe_r_pos_L_des);
        resHand = kinDynSolver->computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des);
        Eigen::VectorXd qIniDes = Eigen::VectorXd::Zero(mj_model->nq, 1);
        qIniDes.block(7, 0, mj_model->nq - 7, 1) = resLeg.jointPosRes + resHand.jointPosRes;
        WBC_solv->setQini(qIniDes, RobotState->q);

        MPC_count = 0; // count for controlling the mpc running period

        startSteppingTime = 3;
        startWalkingTime = 5;
        simEndTime = 30;

        simstart = mj_data->time;
        simTime = mj_data->time;

        m_thread = std::make_unique<std::thread>(
            [this]
            {
                auto startTime = std::chrono::steady_clock::now();
                while (m_running)
                {
                    auto now = std::chrono::steady_clock::now();
                    if (std::chrono::duration_cast<std::chrono::milliseconds>(now - startTime).count() > 16)
                    {
                        std::this_thread::sleep_for(std::chrono::milliseconds(1));
                        startTime = now;
                    }

                    std::lock_guard<std::mutex> guard(m_mutex);
                    Loop();
                }
            });
    }

    void Lock()
    {
        m_mutex.lock();
    }

    void Unlock()
    {
        m_mutex.unlock();
    }

    void Loop()
    {
        mj_step(mj_model, mj_data);

        simTime = mj_data->time;
        // Read the sensors:
        mj_interface->updateSensorValues();
        mj_interface->dataBusWrite(*RobotState);

        // update kinematics and dynamics info
        kinDynSolver->dataBusRead(*RobotState);
        kinDynSolver->computeJ_dJ();
        kinDynSolver->computeDyn();
        kinDynSolver->dataBusWrite(*RobotState);

        // joint number: arm-l: 0-6, arm-r: 7-13, head: 14, waist: 15-17, leg-l: 18-23, leg-r: 24-29

        if (simTime > startWalkingTime)
        {
            jsInterp->setWzDesLPara(0, 1);
            jsInterp->setVxDesLPara(xv_des, 2.0); // jsInterp->setVxDesLPara(0.9,1);
            RobotState->motionState = DataBus::Walk; // start walking
        }
        else
            jsInterp->setIniPos(RobotState->q(0), RobotState->q(1), RobotState->base_rpy(2));
        jsInterp->step();
        RobotState->js_pos_des(2) = stand_legLength + foot_height; // pos z is not assigned in jyInterp
        jsInterp->dataBusWrite(*RobotState); // only pos x, pos y, theta z, vel x, vel y , omega z are rewrote.

        if (simTime >= startSteppingTime)
        {
            // gait scheduler
            gaitScheduler->dataBusRead(*RobotState);
            gaitScheduler->step();
            gaitScheduler->dataBusWrite(*RobotState);

            footPlacement->dataBusRead(*RobotState);
            footPlacement->getSwingPos();
            footPlacement->dataBusWrite(*RobotState);
        }

        // ------------- MPC ------------
        MPC_count = MPC_count + 1;
        if (MPC_count > (dt_200Hz / dt - 1))
        {
            MPC_solv->dataBusRead(*RobotState);
            MPC_solv->cal();
            MPC_solv->dataBusWrite(*RobotState);
            MPC_count = 0;
        }

        // ------------- WBC ------------
        // WBC Calculation
        WBC_solv->dataBusRead(*RobotState);
        WBC_solv->computeDdq(*kinDynSolver);
        WBC_solv->computeTau();
        WBC_solv->dataBusWrite(*RobotState);
        // get the final joint command
        if (simTime <= startSteppingTime)
        {
            RobotState->motors_pos_des = eigen2std(resLeg.jointPosRes + resHand.jointPosRes);
            RobotState->motors_vel_des = motors_vel_des;
            RobotState->motors_tor_des = motors_tau_des;
        }
        else
        {
            MPC_solv->enable();
            Eigen::Matrix<double, 1, nx> L_diag;
            Eigen::Matrix<double, 1, nu> K_diag;
            L_diag << 1.0, 1.0, 1.0, // eul
                1.0, 200.0, 1.0, // pCoM
                1e-7, 1e-7, 1e-7, // w
                100.0, 10.0, 1.0; // vCoM
            K_diag << 1.0, 1.0, 1.0, // fl
                1.0, 1.0, 1.0, 1.0, 1.0, 1.0, // fr
                1.0, 1.0, 1.0, 1.0;
            MPC_solv->set_weight(1e-6, L_diag, K_diag);

            Eigen::VectorXd pos_des = kinDynSolver->integrateDIY(RobotState->q, RobotState->wbc_delta_q_final);
            RobotState->motors_pos_des = eigen2std(pos_des.block(7, 0, model_nv - 6, 1));
            RobotState->motors_vel_des = eigen2std(RobotState->wbc_dq_final);
            RobotState->motors_tor_des = eigen2std(RobotState->wbc_tauJointRes);
        }

        // joint PVT controller
        pvtCtr->dataBusRead(*RobotState);
        if (simTime <= 3)
        {
            pvtCtr->calMotorsPVT(100.0 / 1000.0 / 180.0 * 3.1415);
        }
        else
        {
            pvtCtr->setJointPD(100, 10, "J_ankle_l_pitch");
            pvtCtr->setJointPD(100, 10, "J_ankle_l_roll");
            pvtCtr->setJointPD(100, 10, "J_ankle_r_pitch");
            pvtCtr->setJointPD(100, 10, "J_ankle_r_roll");
            pvtCtr->setJointPD(1000, 100, "J_knee_l_pitch");
            pvtCtr->setJointPD(1000, 100, "J_knee_r_pitch");
            pvtCtr->calMotorsPVT();
        }
        pvtCtr->dataBusWrite(*RobotState);

        // give the joint torque command to Webots
        mj_interface->setMotorsTorque(RobotState->motors_tor_out);
    }

private:
    const double dt = 0.001;
    const double dt_200Hz = 0.005;

    mjModel* mj_model;
    mjData* mj_data;

    int MPC_count = 0; // count for controlling the mpc running period

    double startSteppingTime = 3;
    double startWalkingTime = 5;
    double simEndTime = 30;

    mjtNum simstart;
    double simTime;

    std::unique_ptr<MJ_Interface> mj_interface;
    std::unique_ptr<Pin_KinDyn> kinDynSolver;
    std::unique_ptr<DataBus> RobotState;
    std::unique_ptr<JoyStickInterpreter> jsInterp;
    std::unique_ptr<PVT_Ctr> pvtCtr;
    std::unique_ptr<GaitScheduler> gaitScheduler;
    std::unique_ptr<FootPlacement> footPlacement;
    std::unique_ptr<WBC_priority> WBC_solv;
    std::unique_ptr<MPC> MPC_solv;

    // initialize variables
    double stand_legLength = 1.01; // 0.97;// desired baselink height
    double foot_height = 0.07; // distance between the foot ankel joint and the bottom
    double xv_des = 1.2; // desired velocity in x direction

    Pin_KinDyn::IkRes resLeg;
    Pin_KinDyn::IkRes resHand;

    std::vector<double> motors_pos_des;
    std::vector<double> motors_pos_cur;
    std::vector<double> motors_vel_des;
    std::vector<double> motors_vel_cur;
    std::vector<double> motors_tau_des;
    std::vector<double> motors_tau_cur;

    int model_nv;

    std::unique_ptr<std::thread> m_thread;
    std::mutex m_mutex;

    bool m_running = true;
};


extern "C" __declspec(dllexport) void* Init(mjModel* model, mjData* data)
{
    auto imp = new WalkMpcWbcControllerImp();
    imp->Init(model, data);
    return imp;
}

extern "C" __declspec(dllexport) void Uninit(void* imp)
{
    WalkMpcWbcControllerImp* p = (WalkMpcWbcControllerImp*)imp;
    delete p;
}

extern "C" __declspec(dllexport) void Lock(void* imp)
{
    ((WalkMpcWbcControllerImp*)imp)->Lock();
}

extern "C" __declspec(dllexport) void Unlock(void* imp)
{
    ((WalkMpcWbcControllerImp*)imp)->Unlock();
}
